/*
 * Copyright 2014 The Closure Compiler Authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.google.javascript.jscomp;

import static com.google.common.base.Preconditions.checkState;
import static com.google.common.truth.Truth.assertThat;
import static com.google.javascript.rhino.testing.NodeSubject.assertNode;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
import com.google.javascript.jscomp.CompilerOptions.LanguageMode;
import com.google.javascript.jscomp.colors.Color;
import com.google.javascript.jscomp.colors.StandardColors;
import com.google.javascript.jscomp.testing.CodeSubTree;
import com.google.javascript.jscomp.testing.TestExternsBuilder;
import com.google.javascript.rhino.Node;
import com.google.javascript.rhino.testing.NodeSubject;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

/** Tests for {@link Es6RewriteGenerators} */
@RunWith(JUnit4.class)
public final class Es6RewriteGeneratorsTest extends CompilerTestCase {

  /**
   * Map from special variable names to placeholders.
   *
   * <p>`Es6RewriteGenerators` has to generate unique variable names using
   * `compiler.getUniqueIdSupplier().getUniqueId(compilerInput)`. In order to test this, we will use
   * a utility function that replaces special strings in the expected code with a combination of a
   * prefix and a compiler-input-specific hash.
   *
   * <p>This map defines the special strings that should appear in the test cases. Note that the
   * actual variables generated by `Es6RewriteGenerators` will include an additional "$" + INDEX
   * string on the end. Unless there's more than one generator function in a test case, INDEX will
   * be "0", so you would refer to the context value in the expected output code as "GEN_CONTEXT$0".
   * Each generator function rewritten within a single test case will get a unique INDEX in the
   * order that they are visited.
   */
  private static final ImmutableMap<String, String> SPECIAL_VARIABLES_MAP =
      ImmutableMap.of(
          "GEN_ARGUMENTS", "$jscomp$generator$arguments$",
          "GEN_CONTEXT", "$jscomp$generator$context$",
          "GEN_FORIN", "$jscomp$generator$forin$",
          "GEN_FUNC", "$jscomp$generator$function$",
          "GEN_THIS", "$jscomp$generator$this$");

  public Es6RewriteGeneratorsTest() {
    super(
        new TestExternsBuilder()
            .addAsyncIterable()
            .addArray()
            .addArguments()
            .addObject()
            .addMath()
            .addExtra(
                // stubs of runtime libraries
                """
                /** @const */
                var $jscomp = {};
                $jscomp.generator = {};
                $jscomp.generator.createGenerator = function() {};
                /** @constructor */
                $jscomp.generator.Context = function() {};
                /** @constructor */
                $jscomp.generator.Context.PropertyIterator = function() {};
                $jscomp.asyncExecutePromiseGeneratorFunction = function(program) {};
                $jscomp.asyncExecutePromiseGeneratorProgram = function(program) {};
                """)
            .build());
  }

  @Override
  @Before
  public void setUp() throws Exception {
    super.setUp();
    setAcceptedLanguage(LanguageMode.ECMASCRIPT_2015);
    enableNormalize();
    enableTypeCheck();
    enableTypeInfoValidation();
    replaceTypesWithColors();
    enableMultistageCompilation();
    setGenericNameReplacements(SPECIAL_VARIABLES_MAP);
  }

  @Override
  protected CompilerOptions getOptions() {
    CompilerOptions options = super.getOptions();
    options.setLanguageOut(LanguageMode.ECMASCRIPT3);
    return options;
  }

  @Override
  protected CompilerPass getProcessor(final Compiler compiler) {
    return new Es6RewriteGenerators(compiler);
  }

  private void rewriteGeneratorBody(String beforeBody, String afterBody) {
    rewriteGeneratorBodyWithVars(beforeBody, "", afterBody);
  }

  /** Verifies that generator functions are rewritten to a state machine program. */
  private void rewriteGeneratorBodyWithVars(String beforeBody, String varDecls, String afterBody) {
    rewriteGeneratorBodyWithVarsAndReturnType(beforeBody, varDecls, afterBody, "?");
  }

  private void rewriteGeneratorBodyWithVarsAndReturnType(
      String beforeBody, String varDecls, String afterBody, String returnType) {
    Sources srcs =
        srcs(
            """
            /** @return {RETURN_TYPE} */
            function *genTestFunction() {
            BEFORE_BODY
            }
            """
                .replace("RETURN_TYPE", returnType)
                .replace("BEFORE_BODY", beforeBody));
    Expected originalExpected =
        expected(
            """
            function genTestFunction() {
            VAR_DECLS
              return $jscomp.generator.createGenerator(
                  genTestFunction,
                  function (GEN_CONTEXT$0) {
            AFTER_BODY
                  });
            }
            """
                .replace("VAR_DECLS", varDecls)
                .replace("AFTER_BODY", afterBody));
    test(srcs, originalExpected);
  }

  private void rewriteGeneratorSwitchBody(String beforeBody, String afterBody) {
    rewriteGeneratorSwitchBodyWithVars(beforeBody, "", afterBody);
  }

  /**
   * Verifies that generator functions are rewriteen to a state machine program contining {@code
   * switch} statement.
   *
   * <p>This is the case when total number of program states is more than 3.
   */
  private void rewriteGeneratorSwitchBodyWithVars(
      String beforeBody, String varDecls, String afterBody) {
    rewriteGeneratorBodyWithVars(
        beforeBody,
        varDecls,
        """
        switch (GEN_CONTEXT$0.nextAddress) {
          case 1:
        AFTER_BODY
        }
        """
            .replace("AFTER_BODY", afterBody));
  }

  @Test
  public void testGeneratorForAsyncFunction() {
    test(
        """
        f = function() {
          return (0, $jscomp.asyncExecutePromiseGeneratorFunction)(
              function *() {
                var x = 6;
                yield x;
              });
        }
        """,
        """
        f = function () {
          var x;
          return (0, $jscomp.asyncExecutePromiseGeneratorProgram)(
              function (GEN_CONTEXT$0) {
                 x = 6;
                 return GEN_CONTEXT$0.yield(x, 0);
              });
        }
        """);

    test(
        """
        f = function(x) {
          var $jscomp$restParams = [];
          for (var $jscomp$restIndex = 0;
              $jscomp$restIndex < arguments.length;
              ++$jscomp$restIndex) {
            $jscomp$restParams[$jscomp$restIndex - 0] = arguments[$jscomp$restIndex];
          }
          {
            var bla$0 = $jscomp$restParams;
            return (0, $jscomp.asyncExecutePromiseGeneratorFunction)(
                function *() {
                  var y = bla$0[0];
                  yield y;
                });
          }
        }
        """,
        """
        f = function (x) {
          var $jscomp$restParams = [];
          var $jscomp$restIndex = 0;
          for (;
              $jscomp$restIndex < arguments.length;
              ++$jscomp$restIndex) {
            $jscomp$restParams[$jscomp$restIndex - 0] = arguments[$jscomp$restIndex];
          }
          {
            var bla$0 = $jscomp$restParams;
            var y;
            return (0, $jscomp.asyncExecutePromiseGeneratorProgram)(
                function (GEN_CONTEXT$0) {
                  y = bla$0[0];
                  return GEN_CONTEXT$0.yield(y, 0);
                });
          }
        }
        """);
  }

  @Test
  public void testUnnamed() {
    test(
        "f = function *() {};",
        """
        f = function GEN_FUNC$0() {
          return $jscomp.generator.createGenerator(
              GEN_FUNC$0,
              function (GEN_CONTEXT$0) {
                 GEN_CONTEXT$0.jumpToEnd();
              });
        }
        """);
  }

  @Test
  public void testConst() {
    test(
        """
        /** @const */
        var f = function *() {};
        """,
        """
        /** @const */
        var f = function GEN_FUNC$0() {
          return $jscomp.generator.createGenerator(
              GEN_FUNC$0,
              function (GEN_CONTEXT$0) {
                 GEN_CONTEXT$0.jumpToEnd();
              });
        }
        """);
  }

  @Test
  public void testContainsClosure() {
    test(
        """
        function *f(o) {
          var i = 0
          var pp;
          function msg(p) {
            return i++ + ' ' + p + '\\n';
          }
          for (pp in obj) {
            yield msg(pp);
          }
        }
        """,
        """
        function f(o) {
          function msg(p) {
            return i++ + ' ' + p + '\\n';
          }
          var i
          var pp;
          var GEN_FORIN$0$0;
          return $jscomp.generator.createGenerator(
              f,
              function (GEN_CONTEXT$0) {
                 if (GEN_CONTEXT$0.nextAddress == 1) {
                   i = 0;
                   GEN_FORIN$0$0 = GEN_CONTEXT$0.forIn(obj);
                 }
                 if (!((pp = GEN_FORIN$0$0.getNext()) != null)) {
                   return GEN_CONTEXT$0.jumpTo(0);
                 }
                 return GEN_CONTEXT$0.yield(msg(pp), 2);
              });
        }
        """);
  }

  @Test
  public void testContainsClosureAssignedToVariable() {
    test(
        """
        function *f(o) {
          var i = 0
          var pp;
          var msg = function(p) {
            return i++ + ' ' + p + '\\n';
          }
          for (pp in obj) {
            yield msg(pp);
          }
        }
        """,
        """
        function f(o) {
          var i
          var pp;
          var msg;
          var GEN_FORIN$0$0;
          return $jscomp.generator.createGenerator(
              f,
              function (GEN_CONTEXT$0) {
                 if (GEN_CONTEXT$0.nextAddress == 1) {
                   i = 0;
        // TODO(bradfordcsmith): Maybe it would be better if we hoisted this function
        // expression out so it doesn't create a new closure every time we enter this
        // callback function?
                   msg = function(p) {
                     return i++ + ' ' + p + '\\n';
                   };
                   GEN_FORIN$0$0 = GEN_CONTEXT$0.forIn(obj);
                 }
                 if (!((pp = GEN_FORIN$0$0.getNext()) != null)) {
                   return GEN_CONTEXT$0.jumpTo(0);
                 }
                 return GEN_CONTEXT$0.yield(msg(pp), 2);
              });
        }
        """);
  }

  @Test
  public void testContainsClosureAndTranspiledFromAsync() {
    test(
        """
        function f(o) {
          use(o);
          return (0, $jscomp.asyncExecutePromiseGeneratorFunction)(
              function *() {
                var i = 0
                var pp;
                function msg(p) {
                  return i++ + ' ' + p + '\\n';
                }
                for (pp in obj) {
                  yield msg(pp);
                }
              });
        }
        """,
        """
        function f(o) {
        // When Es6RewriteGenerators recognizes that it is transpiling code that was
        // originally an async function, it will use the body of that function as the
        // place to move variable declarations instead of putting them in the generator
        // function body itself.
        // To maintain normalization, all function declarations must go to the top
          function msg(p) {
            return i++ + ' ' + p + '\\n';
          }
        // In a real situation the only thing that is likely to be here is code that
        // was generated by transpilations that occurred after async function transpilation
        // and before generator transpilation. For example code for handling parameter
        // destructuring.
          use(o);
        // Regular var declarations just go in the order they originally appeared
        // immediately before the executor call.
          var i
          var pp;
          var GEN_FORIN$0$0;
          return (0, $jscomp.asyncExecutePromiseGeneratorProgram)(
              function (GEN_CONTEXT$0) {
                 if (GEN_CONTEXT$0.nextAddress == 1) {
                   i = 0;
                   GEN_FORIN$0$0 = GEN_CONTEXT$0.forIn(obj);
                 }
                 if (!((pp = GEN_FORIN$0$0.getNext()) != null)) {
                   return GEN_CONTEXT$0.jumpTo(0);
                 }
                 return GEN_CONTEXT$0.yield(msg(pp), 2);
              });
        }
        """);
  }

  @Test
  public void testSimpleGenerator() {
    rewriteGeneratorBody("", "  GEN_CONTEXT$0.jumpToEnd();");

    rewriteGeneratorBody("yield;", "  return GEN_CONTEXT$0.yield(void 0, 0);");

    rewriteGeneratorBody("yield 1;", "  return GEN_CONTEXT$0.yield(1, 0);");

    test(
        "/** @param {*} x */ function *f(x, y) {}",
        """
        function f(x, y) {
          return $jscomp.generator.createGenerator(
              f,
              function(GEN_CONTEXT$0) {
                GEN_CONTEXT$0.jumpToEnd();
              });
        }
        """);

    rewriteGeneratorBodyWithVars(
        "var i = 0, j = 2;",
        """
        var i;
        var j;
        """,
        """
        i = 0;
        j = 2;
        GEN_CONTEXT$0.jumpToEnd();
        """);

    rewriteGeneratorBodyWithVars(
        "var i = 0; yield i; i = 1; yield i; i = i + 1; yield i;",
        "var i;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          i = 0;
          return GEN_CONTEXT$0.yield(i, 2);
        }
        if (GEN_CONTEXT$0.nextAddress != 3) {
          i = 1;
          return GEN_CONTEXT$0.yield(i, 3);
        }
        i = i + 1;
        return GEN_CONTEXT$0.yield(i, 0);
        """);
  }

  @Test
  public void testForLoopWithExtraVarDeclaration() {
    test(
        """
        function *gen() {
          var i = 2;
          yield i;
          use(i);
          for (var i = 0; i < 3; i++) {
            use(i);
          }
        }
        """,
        """
        function gen(){
          var i;
          return $jscomp.generator.createGenerator(
              gen,
              function(GEN_CONTEXT$0) {
                if (GEN_CONTEXT$0.nextAddress==1) {
                  i=2;
                  return GEN_CONTEXT$0.yield(i,2);
                }
                use(i);
                i = 0;
                for (; i < 3 ; i++) use(i);
                GEN_CONTEXT$0.jumpToEnd();
              })
        }
        """);
  }

  @Test
  public void testUnreachableCodeGeneration() {
    rewriteGeneratorBody(
        "if (i) return 1; else return 2;",
        """
          if (i) {
            return GEN_CONTEXT$0.return(1);
          } else {
            return GEN_CONTEXT$0.return(2);
          }
        // TODO(b/73762053): Avoid generating unreachable statements.
          GEN_CONTEXT$0.jumpToEnd();
        """);
  }

  @Test
  public void testReturnGenerator() {
    test(
        "function f() { return function *g() {yield 1;} }",
        """
        function f() {
          return function g() {
            return $jscomp.generator.createGenerator(
                g,
                function(GEN_CONTEXT$0) {
                  return GEN_CONTEXT$0.yield(1, 0);
                });
          }
        }
        """);
  }

  @Test
  public void testNestedGenerator() {
    test(
        "function *f() { function *g() {yield 2;} yield 1; }",
        """
        function f() {
          function g() {
            return $jscomp.generator.createGenerator(
                g,
                function(GEN_CONTEXT$0) {
                  return GEN_CONTEXT$0.yield(2, 0);
                });
          }
          return $jscomp.generator.createGenerator(
              f,
              function(GEN_CONTEXT$1) {
                return GEN_CONTEXT$1.yield(1, 0);
              });
        }
        """);
  }

  @Test
  public void testForLoops() {
    rewriteGeneratorBodyWithVars(
        "var i = 0; for (var j = 0; j < 10; j++) { i += j; }",
        "var i; var j;",
        """
        i = 0;
        j = 0;
        for (; j < 10; j++) {
          i = i + j; // normalization rewrote this
        }
        GEN_CONTEXT$0.jumpToEnd();
        """);

    rewriteGeneratorBodyWithVars(
        "var i = 0; for (var j = yield; j < 10; j++) { i += j; }",
        "var i; var j;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          i = 0;
          return GEN_CONTEXT$0.yield(void 0, 2);
        }
        j = GEN_CONTEXT$0.yieldResult;
        for (; j < 10; j++) {
          i = i + j; // normalization rewrote this
        }
        GEN_CONTEXT$0.jumpToEnd();
        """);

    rewriteGeneratorBody("for (;;) { yield 1; }", "  return GEN_CONTEXT$0.yield(1, 1);");

    rewriteGeneratorBodyWithVars(
        "for (var yieldResult; yieldResult === undefined; yieldResult = yield 1) {}",
        "var yieldResult;",
        """
          if (GEN_CONTEXT$0.nextAddress == 1) {
            if (!(yieldResult === undefined)) return GEN_CONTEXT$0.jumpTo(0);
            return GEN_CONTEXT$0.yield(1,5);
          }
          yieldResult = GEN_CONTEXT$0.yieldResult;
          return GEN_CONTEXT$0.jumpTo(1);
        """);

    rewriteGeneratorBodyWithVars(
        "for (var j = 0; j < 10; j++) { yield j; }",
        "var j;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          j = 0;
        }
        if (GEN_CONTEXT$0.nextAddress != 3) {
          if (!(j < 10)) {
            return GEN_CONTEXT$0.jumpTo(0);
          }
          return GEN_CONTEXT$0.yield(j, 3);
        }
        j++;
        return GEN_CONTEXT$0.jumpTo(2);
        """);

    rewriteGeneratorBodyWithVars(
        "var i = 0; for (var j = 0; j < 10; j++) { i += j; yield 5; }",
        "var i; var j;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          i = 0;
          j = 0;
        }
        if (GEN_CONTEXT$0.nextAddress != 3) {
          if (!(j < 10)) {
            return GEN_CONTEXT$0.jumpTo(0);
          }
          i = i + j; // normalization rewrote this
          return GEN_CONTEXT$0.yield(5, 3);
        }
        j++;
        return GEN_CONTEXT$0.jumpTo(2);
        """);
  }

  @Test
  public void testWhileLoops() {
    rewriteGeneratorBodyWithVars(
        "var i = 0; while (i < 10) { i++; i++; i++; } yield i;",
        "  var i;",
        """
        i = 0;
        for (; i < 10;) { i ++; i++; i++; }
        return GEN_CONTEXT$0.yield(i, 0);
        """);

    rewriteGeneratorSwitchBodyWithVars(
        "var j = 0; while (j < 10) { yield j; j++; } j += 10;",
        "var j;",
        """
          j = 0;
        case 2:
          if (!(j < 10)) {
            GEN_CONTEXT$0.jumpTo(4);
            break;
          }
          return GEN_CONTEXT$0.yield(j, 5)
        case 5:
          j++;
          GEN_CONTEXT$0.jumpTo(2);
          break;
        case 4:
          j = j + 10; // normalization rewrote this
          GEN_CONTEXT$0.jumpToEnd();
        """);

    rewriteGeneratorBodyWithVars(
        "var j = 0; while (j < 10) { yield j; j++; } yield 5",
        "var j;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          j = 0;
        }
        if (GEN_CONTEXT$0.nextAddress != 5) {
          if (!(j < 10)) {
            return GEN_CONTEXT$0.yield(5, 0);
          }
          return GEN_CONTEXT$0.yield(j, 5)
        }
          j++;
          return GEN_CONTEXT$0.jumpTo(2);
        """);

    rewriteGeneratorBodyWithVars(
        "var j = 0; while (yield) { j++; }",
        "var j;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          j = 0;
        }
        if (GEN_CONTEXT$0.nextAddress != 5) {
          return GEN_CONTEXT$0.yield(void 0, 5);
        }
        if (!(GEN_CONTEXT$0.yieldResult)) {
          return GEN_CONTEXT$0.jumpTo(0);
        }
        j++;
        return GEN_CONTEXT$0.jumpTo(2);
        """);
  }

  @Test
  public void testDecomposeComplexYieldExpression() {
    test(
        """
        /* @return {?} */ function *f() {
          var o = {bar: function(x) {}};
          (yield 5) && o.bar(yield 5);
        }
        """,
        """
        function f(){
          var o;
          var JSCompiler_temp$jscomp$0;
          var JSCompiler_temp_const$jscomp$2;
          var JSCompiler_temp_const$jscomp$1;
          var JSCompiler_temp_const$jscomp$3;
          return $jscomp.generator.createGenerator(f, function(GEN_CONTEXT$0) {
            switch(GEN_CONTEXT$0.nextAddress) {
              case 1:
                o = {bar:function(x) {}};
                return GEN_CONTEXT$0.yield(5, 2);
              case 2:
                if (!(JSCompiler_temp$jscomp$0 = GEN_CONTEXT$0.yieldResult)) {
                  GEN_CONTEXT$0.jumpTo(3);
                  break;
                }
                JSCompiler_temp_const$jscomp$2 = o;
                JSCompiler_temp_const$jscomp$1 = JSCompiler_temp_const$jscomp$2.bar;
                JSCompiler_temp_const$jscomp$3 = JSCompiler_temp_const$jscomp$2;
                return GEN_CONTEXT$0.yield(5, 4);
              case 4:
                JSCompiler_temp$jscomp$0 =
                    JSCompiler_temp_const$jscomp$1.call(
                        JSCompiler_temp_const$jscomp$3,
                        GEN_CONTEXT$0.yieldResult);
              case 3:
                JSCompiler_temp$jscomp$0;
                GEN_CONTEXT$0.jumpToEnd();
            }
          });
        }
        """);
  }

  @Test
  public void testDecomposableExpression() {
    rewriteGeneratorBodyWithVars(
        "return a + (a = b) + (b = yield) + a;",
        "var JSCompiler_temp_const$jscomp$0;",
        """
          if (GEN_CONTEXT$0.nextAddress == 1) {
            JSCompiler_temp_const$jscomp$0 = a + (a = b);
            return GEN_CONTEXT$0.yield(void 0, 2);
          }
          return GEN_CONTEXT$0.return(
        JSCompiler_temp_const$jscomp$0 + (b = GEN_CONTEXT$0.yieldResult) + a);
        """);

    rewriteGeneratorSwitchBodyWithVars(
        "return (yield ((yield 1) + (yield 2)));",
        "var JSCompiler_temp_const$jscomp$0;",
        """
          return GEN_CONTEXT$0.yield(1, 3);
        case 3:
          JSCompiler_temp_const$jscomp$0=GEN_CONTEXT$0.yieldResult;
          return GEN_CONTEXT$0.yield(2, 4);
        case 4:
          return GEN_CONTEXT$0.yield(
              JSCompiler_temp_const$jscomp$0 + GEN_CONTEXT$0.yieldResult, 2);
        case 2:
          return GEN_CONTEXT$0.return(GEN_CONTEXT$0.yieldResult);
        """);

    rewriteGeneratorBodyWithVars(
        "var o = {bar: function(x) {}}; o.bar(yield 5);",
        """
        var o;
        var JSCompiler_temp_const$jscomp$1;
        var JSCompiler_temp_const$jscomp$0;
        """,
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          o = {bar: function(x) {}};
          JSCompiler_temp_const$jscomp$1 = o;
          JSCompiler_temp_const$jscomp$0 = JSCompiler_temp_const$jscomp$1.bar;
          return GEN_CONTEXT$0.yield(5, 2);
        }
        JSCompiler_temp_const$jscomp$0.call(
            JSCompiler_temp_const$jscomp$1, GEN_CONTEXT$0.yieldResult);
        GEN_CONTEXT$0.jumpToEnd();
        """);
  }

  @Test
  public void testGeneratorCannotConvertYet() {
    testError(
        "function *f(b, i) {switch (i) { case yield: return b; }}",
        TranspilationUtil.CANNOT_CONVERT_YET);
  }

  @Test
  public void testThrow() {
    rewriteGeneratorBody("throw 1;", "throw 1;");
  }

  @Test
  public void testLabels() {
    rewriteGeneratorBody(
        "l: if (true) { break l; }",
        """
          l: {
            if (true) { break l; }
          }
          GEN_CONTEXT$0.jumpToEnd();
        """);

    rewriteGeneratorBody(
        "l: if (yield) { break l; }",
        """
          if (GEN_CONTEXT$0.nextAddress == 1)
            return GEN_CONTEXT$0.yield(void 0, 3);
          if (GEN_CONTEXT$0.yieldResult) {
            return GEN_CONTEXT$0.jumpTo(0);
          }
          GEN_CONTEXT$0.jumpToEnd();
        """);

    rewriteGeneratorBody(
        "l: if (yield) { while (1) {break l;} }",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          return GEN_CONTEXT$0.yield(void 0, 3);
        }
        if (GEN_CONTEXT$0.yieldResult) {
          for (; 1;) {
            return GEN_CONTEXT$0.jumpTo(0);
          }
        }
        GEN_CONTEXT$0.jumpToEnd();
        """);

    rewriteGeneratorBody(
        "l: for (;;) { yield i; continue l; }",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          return GEN_CONTEXT$0.yield(i, 5);
        }
        return GEN_CONTEXT$0.jumpTo(1);
        return GEN_CONTEXT$0.jumpTo(1);
        """);

    rewriteGeneratorBody(
        "l1: l2: if (yield) break l1; else break l2;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          return GEN_CONTEXT$0.yield(void 0, 3);
        }
        if(GEN_CONTEXT$0.yieldResult) {
          return GEN_CONTEXT$0.jumpTo(0);
        } else {
          return GEN_CONTEXT$0.jumpTo(0);
        }
        GEN_CONTEXT$0.jumpToEnd();
        """);
  }

  @Test
  public void testUnreachable() {
    // TODO(skill): The generator transpilation should not produce any unreachable code
    rewriteGeneratorBody(
        "while (true) {yield; break;}",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          if (!true) {
            return GEN_CONTEXT$0.jumpTo(0);
          }
          return GEN_CONTEXT$0.yield(void 0, 5);
        }
        return GEN_CONTEXT$0.jumpTo(0);
        return GEN_CONTEXT$0.jumpTo(1);
        """);
  }

  @Test
  public void testCaseNumberOptimization() {
    rewriteGeneratorBodyWithVars(
        """
        for (;;) {
          var gen = generatorSrc();
          var gotResponse = false;
          for (var response in gen) {
            yield response;
            gotResponse = true;
          }
          if (!gotResponse) {
            return;
          }
        }
        """,
        """
        var gen;
        var gotResponse;
        var response;
        var GEN_FORIN$0$0;
        """,
        """
        switch (GEN_CONTEXT$0.nextAddress) {
          case 1:
            gen = generatorSrc();
            gotResponse = false;
            GEN_FORIN$0$0 = GEN_CONTEXT$0.forIn(gen);
          case 5:
            if (!((response=GEN_FORIN$0$0.getNext()) != null)) {
              GEN_CONTEXT$0.jumpTo(7);
              break;
            }
            return GEN_CONTEXT$0.yield(response, 8);
          case 8:
            gotResponse = true;
            GEN_CONTEXT$0.jumpTo(5);
            break;
          case 7:
            if (!gotResponse) return GEN_CONTEXT$0.return();
            GEN_CONTEXT$0.jumpTo(1);
            break;
        }
        """);

    rewriteGeneratorBody(
        "do { do { do { yield; } while (3) } while (2) } while (1)",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          return GEN_CONTEXT$0.yield(void 0, 10);
        }
        if (3) {
          return GEN_CONTEXT$0.jumpTo(1);
        }
        if (2) {
          return GEN_CONTEXT$0.jumpTo(1);
        }
        if (1) {
          return GEN_CONTEXT$0.jumpTo(1);
        }
        GEN_CONTEXT$0.jumpToEnd();
        """);
  }

  @Test
  public void testIf() {
    rewriteGeneratorBodyWithVars(
        "var j = 0; if (yield) { j = 1; }",
        "var j;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          j = 0;
          return GEN_CONTEXT$0.yield(void 0, 2);
        }
        if (GEN_CONTEXT$0.yieldResult) { j = 1; }
        GEN_CONTEXT$0.jumpToEnd();
        """);

    rewriteGeneratorBodyWithVars(
        "var j = 0; if (j < 1) { j = 5; } else { yield j; }",
        "var j;",
        """
        j = 0;
        if (j < 1) {
          j = 5;
          return GEN_CONTEXT$0.jumpTo(0);
        }
        return GEN_CONTEXT$0.yield(j, 0);
        """);

    // When "else" doesn't contain yields, it's more optimal to swap "if" and else "blocks" and
    // negate the condition.
    rewriteGeneratorBodyWithVars(
        "var j = 0; if (j < 1) { yield j; } else { j = 5; }",
        "var j;",
        """
        j = 0;
        if (!(j < 1)) {
          j = 5;
          return GEN_CONTEXT$0.jumpTo(0);
        }
        return GEN_CONTEXT$0.yield(j, 0);
        """);

    // No "else" block, pretend as it's empty
    rewriteGeneratorBodyWithVars(
        "var j = 0; if (j < 1) { yield j; }",
        "var j;",
        """
        j = 0;
        if (!(j < 1)) {
          return GEN_CONTEXT$0.jumpTo(0);
        }
        return GEN_CONTEXT$0.yield(j, 0);
        """);

    rewriteGeneratorBody(
        "if (i < 1) { yield i; } else { yield 1; }",
        """
        if (i < 1) {
          return GEN_CONTEXT$0.yield(i, 0);
        }
        return GEN_CONTEXT$0.yield(1, 0);
        """);

    rewriteGeneratorSwitchBody(
        "if (i < 1) { yield i; yield i + 1; i = 10; } else { yield 1; yield 2; i = 5;}",
        """
          if (i < 1) {
            return GEN_CONTEXT$0.yield(i, 6);
          }
          return GEN_CONTEXT$0.yield(1, 4);
        case 4:
          return GEN_CONTEXT$0.yield(2, 5);
        case 5:
          i = 5;
          GEN_CONTEXT$0.jumpTo(0);
          break;
        case 6:
          return GEN_CONTEXT$0.yield(i + 1, 7)
        case 7:
          i = 10;
          GEN_CONTEXT$0.jumpToEnd();
        """);

    rewriteGeneratorBody(
        """
        if (i < 1) {
          while (true) {
            yield 1;
          }
        } else {
          while (false) {
            yield 2;
          }
        }
        """,
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          if (i < 1) {
            return GEN_CONTEXT$0.jumpTo(8);
          }
        }
        if (GEN_CONTEXT$0.nextAddress != 8) {
          if (!false) {
            return GEN_CONTEXT$0.jumpTo(0);
          }
          return GEN_CONTEXT$0.yield(2, 4);
        }
        if (!true) {
          return GEN_CONTEXT$0.jumpTo(0);
        }
        return GEN_CONTEXT$0.yield(1, 8);
        """);

    rewriteGeneratorBody(
        "if (i < 1) { if (i < 2) {yield 2; } } else { yield 1; }",
        """
        if (i < 1) {
          if (!(i < 2)) {
            return GEN_CONTEXT$0.jumpTo(0);
          }
          return GEN_CONTEXT$0.yield(2, 0);
        }
        return GEN_CONTEXT$0.yield(1, 0)
        """);

    rewriteGeneratorBody(
        "if (i < 1) { if (i < 2) {yield 2; } else { yield 3; } } else { yield 1; }",
        """
        if (i < 1) {
          if (i < 2) {
            return GEN_CONTEXT$0.yield(2, 0);
          }
          return GEN_CONTEXT$0.yield(3, 0);
        }
        return GEN_CONTEXT$0.yield(1, 0)
        """);
  }

  @Test
  public void testReturn() {
    rewriteGeneratorBody("return 1;", "return GEN_CONTEXT$0.return(1);");

    rewriteGeneratorBodyWithVars(
        "return this;", "var GEN_THIS$0 = this;", "return GEN_CONTEXT$0.return(GEN_THIS$0);");

    rewriteGeneratorBodyWithVars(
        "return this.test({value: this});",
        "var GEN_THIS$0 = this;",
        """
        return GEN_CONTEXT$0.return(
            GEN_THIS$0.test({value: GEN_THIS$0}));
        """);

    rewriteGeneratorBodyWithVars(
        "return this[yield];",
        "var GEN_THIS$0 = this;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1)
          return GEN_CONTEXT$0.yield(void 0, 2);
        return GEN_CONTEXT$0.return(
            GEN_THIS$0[GEN_CONTEXT$0.yieldResult]);
        """);
  }

  @Test
  public void testBreakContinue() {
    rewriteGeneratorBodyWithVars(
        "var j = 0; while (j < 10) { yield j; break; }",
        "var j;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          j = 0;
        }
        if (GEN_CONTEXT$0.nextAddress != 5) {
          if (!(j < 10)) {
            return GEN_CONTEXT$0.jumpTo(0);
          }
          return GEN_CONTEXT$0.yield(j, 5);
        }
        return GEN_CONTEXT$0.jumpTo(0);
        return GEN_CONTEXT$0.jumpTo(2)
        """);

    rewriteGeneratorBodyWithVars(
        "var j = 0; while (j < 10) { yield j; continue; }",
        "var j;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          j = 0;
        }
        if (GEN_CONTEXT$0.nextAddress != 5) {
          if (!(j < 10)) {
            return GEN_CONTEXT$0.jumpTo(0);
          }
          return GEN_CONTEXT$0.yield(j, 5);
        }
        return GEN_CONTEXT$0.jumpTo(2);
        return GEN_CONTEXT$0.jumpTo(2)
        """);

    rewriteGeneratorBodyWithVars(
        "for (var j = 0; j < 10; j++) { yield j; break; }",
        "var j;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          j = 0;
        }
        if (GEN_CONTEXT$0.nextAddress != 5) {
          if (!(j < 10)) {
            return GEN_CONTEXT$0.jumpTo(0);
          }
          return GEN_CONTEXT$0.yield(j, 5);
        }
        return GEN_CONTEXT$0.jumpTo(0);
        j++;
        return GEN_CONTEXT$0.jumpTo(2)
        """);

    rewriteGeneratorSwitchBodyWithVars(
        "for (var j = 0; j < 10; j++) { yield j; continue; }",
        "var j;",
        """
          j = 0;
        case 2:
          if (!(j < 10)) {
            GEN_CONTEXT$0.jumpTo(0);
            break;
          }
          return GEN_CONTEXT$0.yield(j, 5);
        case 5:
          GEN_CONTEXT$0.jumpTo(3);
          break;
        case 3:
          j++;
          GEN_CONTEXT$0.jumpTo(2)
          break;
        """);
  }

  @Test
  public void testDoWhileLoops() {
    rewriteGeneratorBody(
        "do { yield j; } while (j < 10);",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          return GEN_CONTEXT$0.yield(j, 4);
        }
        if (j<10) {
          return GEN_CONTEXT$0.jumpTo(1);
        }
        GEN_CONTEXT$0.jumpToEnd();
        """);

    rewriteGeneratorBody(
        "do {} while (yield 1);",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          return GEN_CONTEXT$0.yield(1, 5);
        }
        if (GEN_CONTEXT$0.yieldResult) {
          return GEN_CONTEXT$0.jumpTo(1);
        }
        GEN_CONTEXT$0.jumpToEnd();
        """);
  }

  @Test
  public void testYieldNoValue() {
    rewriteGeneratorBody("yield;", "return GEN_CONTEXT$0.yield(void 0, 0);");
  }

  @Test
  public void testReturnNoValue() {
    rewriteGeneratorBody(
        "return;", //
        "return GEN_CONTEXT$0.return();");
  }

  @Test
  public void testYieldExpression() {
    rewriteGeneratorBody(
        "return (yield 1);",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          return GEN_CONTEXT$0.yield(1, 2);
        }
        return GEN_CONTEXT$0.return(GEN_CONTEXT$0.yieldResult);
        """);
  }

  @Test
  public void testFunctionInGenerator() {
    rewriteGeneratorBodyWithVars(
        "function g() {}", // body statement is just s function definition
        "function g() {}", // and it remains the only variable declaration after transpilation
        "  GEN_CONTEXT$0.jumpToEnd();");
  }

  @Test
  public void testYieldAll() {
    rewriteGeneratorBody(
        "yield * n;", //
        "  return GEN_CONTEXT$0.yieldAll(n, 0);");

    rewriteGeneratorBodyWithVars(
        "var i = yield * n;",
        "var i;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          return GEN_CONTEXT$0.yieldAll(n, 2);
        }
        i = GEN_CONTEXT$0.yieldResult;
        GEN_CONTEXT$0.jumpToEnd();
        """);
  }

  @Test
  public void testYieldArguments() {
    rewriteGeneratorBodyWithVars(
        "yield arguments[0];",
        "var GEN_ARGUMENTS$0 = arguments;",
        "return GEN_CONTEXT$0.yield(GEN_ARGUMENTS$0[0], 0);");
  }

  @Test
  public void testYieldThis() {
    rewriteGeneratorBodyWithVars(
        "yield this;", "var GEN_THIS$0 = this;", "return GEN_CONTEXT$0.yield(GEN_THIS$0, 0);");
  }

  @Test
  public void testGeneratorShortCircuit() {
    rewriteGeneratorBodyWithVars(
        "0 || (yield 1);",
        "var JSCompiler_temp$jscomp$0;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          if(JSCompiler_temp$jscomp$0 = 0) {
            return GEN_CONTEXT$0.jumpTo(2);
          }
          return GEN_CONTEXT$0.yield(1, 3);
        }
        if (GEN_CONTEXT$0.nextAddress != 2) {
          JSCompiler_temp$jscomp$0=GEN_CONTEXT$0.yieldResult;
        }
        JSCompiler_temp$jscomp$0;
        GEN_CONTEXT$0.jumpToEnd();
        """);

    rewriteGeneratorBodyWithVars(
        "0 && (yield 1);",
        "var JSCompiler_temp$jscomp$0;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          if(!(JSCompiler_temp$jscomp$0=0)) {
            return GEN_CONTEXT$0.jumpTo(2);
          }
          return GEN_CONTEXT$0.yield(1, 3);
        }
        if (GEN_CONTEXT$0.nextAddress != 2) {
          JSCompiler_temp$jscomp$0=GEN_CONTEXT$0.yieldResult;
        }
        JSCompiler_temp$jscomp$0;
        GEN_CONTEXT$0.jumpToEnd();
        """);

    rewriteGeneratorBodyWithVars(
        "0 ? 1 : (yield 1);",
        "var JSCompiler_temp$jscomp$0;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          if(0) {
            JSCompiler_temp$jscomp$0 = 1;
            return GEN_CONTEXT$0.jumpTo(2);
          }
          return GEN_CONTEXT$0.yield(1, 3);
        }
        if (GEN_CONTEXT$0.nextAddress != 2) {
          JSCompiler_temp$jscomp$0 = GEN_CONTEXT$0.yieldResult;
        }
        JSCompiler_temp$jscomp$0;
        GEN_CONTEXT$0.jumpToEnd();
        """);
  }

  @Test
  public void testVar() {
    rewriteGeneratorBodyWithVars(
        "var va = 10, vb, vc = yield 10, vd = yield 20, vf, vg='test';",
        """
        var va;
        var vb;
        var vc;
        var vd;
        var vf;
        var vg;
        """,
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          va = 10;
          return GEN_CONTEXT$0.yield(10, 2);
        }
        if (GEN_CONTEXT$0.nextAddress != 3) {
          vc = GEN_CONTEXT$0.yieldResult;
          return GEN_CONTEXT$0.yield(20, 3);
        }
        vd = GEN_CONTEXT$0.yieldResult;
        vg = 'test';
        GEN_CONTEXT$0.jumpToEnd();
        """);

    rewriteGeneratorBodyWithVars(
        "var /** @const */ va = 10, vb, vc = yield 10, vd = yield 20, vf, vg='test';",
        """
        var /** @const */ va;
        var vb;
        var vc;
        var vd
        var vf;
        var vg;
        """,
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          va = 10;
          return GEN_CONTEXT$0.yield(10, 2);
        }
        if (GEN_CONTEXT$0.nextAddress != 3) {
          vc = GEN_CONTEXT$0.yieldResult;
          return GEN_CONTEXT$0.yield(20, 3);
        }
        vd = GEN_CONTEXT$0.yieldResult;
        vg = 'test';
        GEN_CONTEXT$0.jumpToEnd();
        """);
  }

  @Test
  public void testYieldSwitch() {
    rewriteGeneratorSwitchBody(
        """
        while (1) {
          switch (i) {
            case 1:
              ++i;
              break;
            case 2:
              yield 3;
              continue;
            case 10:
            case 3:
              yield 4;
            case 4:
              return 1;
            case 5:
              return 2;
            default:
              yield 5;
          }
        }
        """,
        """
          if (!1) {
            GEN_CONTEXT$0.jumpTo(0);
            break;
          }
          switch (i) {
            case 1: ++i; break;
            case 2: return GEN_CONTEXT$0.jumpTo(5);
            case 10:
            case 3: return GEN_CONTEXT$0.jumpTo(6);
            case 4: return GEN_CONTEXT$0.jumpTo(7);
            case 5: return GEN_CONTEXT$0.return(2);
            default: return GEN_CONTEXT$0.jumpTo(8);
          }
          GEN_CONTEXT$0.jumpTo(1);
          break;
        case 5: return GEN_CONTEXT$0.yield(3, 10);
        case 10:
          GEN_CONTEXT$0.jumpTo(1);
          break;
        case 6: return GEN_CONTEXT$0.yield(4, 7);
        case 7: return GEN_CONTEXT$0.return(1);
        case 8: return GEN_CONTEXT$0.yield(5, 1);
        """);

    rewriteGeneratorBody(
        """
        switch (yield) {
          default:
          case 1:
            yield 1;}
        """,
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          return GEN_CONTEXT$0.yield(void 0, 2);
        }
        if (GEN_CONTEXT$0.nextAddress != 3) {
          switch (GEN_CONTEXT$0.yieldResult) {
            default:
            case 1:
              return GEN_CONTEXT$0.jumpTo(3)
          }
          return GEN_CONTEXT$0.jumpTo(0);
        }
        return GEN_CONTEXT$0.yield(1, 0);
        """);
  }

  @Test
  public void testNoTranslate() {
    rewriteGeneratorBody(
        "if (1) { try {} catch (e) {} throw 1; }",
        """
        if (1) { try {} catch (e) {} throw 1; }
        GEN_CONTEXT$0.jumpToEnd();
        """);
  }

  @Test
  public void testForIn() {
    rewriteGeneratorBodyWithVars(
        "for (var i in yield) { }",
        "var i;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          return GEN_CONTEXT$0.yield(void 0, 2);
        }
        for (i in GEN_CONTEXT$0.yieldResult) { }
        GEN_CONTEXT$0.jumpToEnd();
        """);

    rewriteGeneratorBodyWithVars(
        "for (var i in j) { yield i; }",
        """
        var i;
        var GEN_FORIN$0$0;
        """,
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          GEN_FORIN$0$0 = GEN_CONTEXT$0.forIn(j);
        }
        if (!((i = GEN_FORIN$0$0.getNext()) != null)) {
          return GEN_CONTEXT$0.jumpTo(0);
        }
        return GEN_CONTEXT$0.yield(i, 2);
        """);

    rewriteGeneratorBodyWithVars(
        "for (var i in yield) { yield i; }",
        """
        var i;
        var GEN_FORIN$0$0;
        """,
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          return GEN_CONTEXT$0.yield(void 0, 2)
        }
        if (GEN_CONTEXT$0.nextAddress != 3) {
          GEN_FORIN$0$0 =
              GEN_CONTEXT$0.forIn(GEN_CONTEXT$0.yieldResult);
        }
        if (!((i = GEN_FORIN$0$0.getNext()) != null)) {
          return GEN_CONTEXT$0.jumpTo(0);
        }
        return GEN_CONTEXT$0.yield(i, 3);
        """);

    rewriteGeneratorBodyWithVars(
        "for (i[yield] in j) {}",
        "var GEN_FORIN$0$0; var JSCompiler_temp_const$jscomp$0;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          GEN_FORIN$0$0 = GEN_CONTEXT$0.forIn(j);
        }
        if (GEN_CONTEXT$0.nextAddress != 5) {
          JSCompiler_temp_const$jscomp$0 = i;
          return GEN_CONTEXT$0.yield(void 0, 5);
        }
        if (!((JSCompiler_temp_const$jscomp$0[GEN_CONTEXT$0.yieldResult] =
            GEN_FORIN$0$0.getNext()) != null)) {
          return GEN_CONTEXT$0.jumpTo(0);
        }
        return GEN_CONTEXT$0.jumpTo(2);
        """);
  }

  @Test
  public void testTryCatch() {
    rewriteGeneratorBodyWithVars(
        "try {yield 1;} catch (e) {}",
        "var e;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          GEN_CONTEXT$0.setCatchFinallyBlocks(2);
          return GEN_CONTEXT$0.yield(1, 4);
        }
        if (GEN_CONTEXT$0.nextAddress != 2) {
          return GEN_CONTEXT$0.leaveTryBlock(0)
        }
        e = GEN_CONTEXT$0.enterCatchBlock();
        GEN_CONTEXT$0.jumpToEnd();
        """);

    rewriteGeneratorSwitchBodyWithVars(
        """
        try {yield 1;} catch (e) { use('first ' + e); }
        try {yield 1;} catch (e) { use('second ' + e)}
        """,
        """
        var e;
        // The second catch variable was renamed by normalization
        var e$jscomp$1;
        """,
        """
          GEN_CONTEXT$0.setCatchFinallyBlocks(2);
          return GEN_CONTEXT$0.yield(1, 4);
        case 4:
          GEN_CONTEXT$0.leaveTryBlock(3)
          break;
        case 2:
          e=GEN_CONTEXT$0.enterCatchBlock();
          use('first ' + e);
        case 3:
          GEN_CONTEXT$0.setCatchFinallyBlocks(5);
          return GEN_CONTEXT$0.yield(1, 7);
        case 7:
          GEN_CONTEXT$0.leaveTryBlock(0)
          break;
        case 5:
          e$jscomp$1 = GEN_CONTEXT$0.enterCatchBlock();
          use('second ' + e$jscomp$1);
          GEN_CONTEXT$0.jumpToEnd();
        """);

    rewriteGeneratorBodyWithVars(
        "l1: try { break l1; } catch (e) { yield; } finally {}",
        "var e;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          GEN_CONTEXT$0.setCatchFinallyBlocks(3, 4);
          return GEN_CONTEXT$0.jumpThroughFinallyBlocks(0);
        }
        if (GEN_CONTEXT$0.nextAddress != 3) {
          GEN_CONTEXT$0.enterFinallyBlock();
          return GEN_CONTEXT$0.leaveFinallyBlock(0);
        }
        e = GEN_CONTEXT$0.enterCatchBlock();
        return GEN_CONTEXT$0.yield(void 0, 4)
        """);
  }

  @Test
  public void testFinally() {
    rewriteGeneratorBodyWithVars(
        "try {yield 1;} catch (e) {} finally {b();}",
        "var e;",
        """
        if (GEN_CONTEXT$0.nextAddress == 1) {
          GEN_CONTEXT$0.setCatchFinallyBlocks(2, 3);
          return GEN_CONTEXT$0.yield(1, 3);
        }
        if (GEN_CONTEXT$0.nextAddress != 2) {
          GEN_CONTEXT$0.enterFinallyBlock();
          b();
          return GEN_CONTEXT$0.leaveFinallyBlock(0);
        }
        e = GEN_CONTEXT$0.enterCatchBlock();
        return GEN_CONTEXT$0.jumpTo(3);
        """);

    rewriteGeneratorSwitchBodyWithVars(
        """
        try {
          try {
            yield 1;
            throw 2;
          } catch (x) {
            throw yield x;
          } finally {
            yield 5;
          }
        } catch (thrown) {
          yield thrown;
        }
        """,
        "var x; var thrown;",
        """
          GEN_CONTEXT$0.setCatchFinallyBlocks(2);
          GEN_CONTEXT$0.setCatchFinallyBlocks(4, 5);
          return GEN_CONTEXT$0.yield(1, 7);
        case 7:
          throw 2;
        case 5:
          GEN_CONTEXT$0.enterFinallyBlock(2);
          return GEN_CONTEXT$0.yield(5, 8);
        case 8:
          GEN_CONTEXT$0.leaveFinallyBlock(6);
          break;
        case 4:
          x = GEN_CONTEXT$0.enterCatchBlock();
          return GEN_CONTEXT$0.yield(x, 9);
        case 9:
          throw GEN_CONTEXT$0.yieldResult;
          GEN_CONTEXT$0.jumpTo(5);
          break;
        case 6:
          GEN_CONTEXT$0.leaveTryBlock(0);
          break;
        case 2:
          thrown = GEN_CONTEXT$0.enterCatchBlock();
          return GEN_CONTEXT$0.yield(thrown,0)
        """);
  }

  /** Tests correctness of type information after transpilation */
  @Test
  public void testYield_withTypes() {
    Node returnNode =
        testAndReturnBodyForNumericGenerator(
            "yield 1 + 2;", "", "return GEN_CONTEXT$0.yield(1 + 2, 0);");

    checkState(returnNode.isReturn(), returnNode);
    Node callNode = returnNode.getFirstChild();
    checkState(callNode.isCall(), callNode);
    // TODO(b/142881197): we can't give a more accurate type right now.
    assertNode(callNode).hasColorThat().isEqualTo(StandardColors.UNKNOWN);

    Node yieldFn = callNode.getFirstChild();
    Node jscompGeneratorContext = yieldFn.getFirstChild();
    Color generatorContext =
        Iterables.getOnlyElement(
            CodeSubTree.findFirstNode(
                    getLastCompiler().getExternsRoot(),
                    (n) -> n.matchesQualifiedName("$jscomp.generator.Context"))
                .getColor()
                .getInstanceColors());
    assertThat(jscompGeneratorContext.getColor()).isEqualTo(generatorContext);

    // Check types on "1 + 2" are still present after transpilation
    Node yieldedValue = callNode.getSecondChild(); // 1 + 2

    checkState(yieldedValue.isAdd(), yieldedValue);
    assertNode(yieldedValue).hasColorThat().isEqualTo(StandardColors.NUMBER);
    assertNode(yieldedValue.getFirstChild()).hasColorThat().isEqualTo(StandardColors.NUMBER); // 1
    assertNode(yieldedValue.getSecondChild()).hasColorThat().isEqualTo(StandardColors.NUMBER); // 2

    Node zero = yieldedValue.getNext();
    checkState(0 == zero.getDouble(), zero);
    assertNode(zero).hasColorThat().isEqualTo(StandardColors.NUMBER);
  }

  @Test
  public void testYieldAll_withTypes() {
    Node returnNode =
        testAndReturnBodyForNumericGenerator(
            "yield * [1, 2];", "", "return GEN_CONTEXT$0.yieldAll([1, 2], 0);");

    checkState(returnNode.isReturn(), returnNode);
    Node callNode = returnNode.getFirstChild();
    checkState(callNode.isCall(), callNode);
    // TODO(b/142881197): we can't give a more accurate type right now.
    assertNode(callNode).hasColorThat().isEqualTo(StandardColors.UNKNOWN);

    Node yieldAllFn = callNode.getFirstChild();
    checkState(yieldAllFn.isGetProp());

    // Check that the original types on "[1, 2]" are still present after transpilation
    Node yieldedValue = callNode.getSecondChild(); // [1, 2]

    checkState(yieldedValue.isArrayLit(), yieldedValue);
    assertNode(yieldedValue)
        .hasColorThat()
        .isEqualTo(getLastCompiler().getColorRegistry().get(StandardColors.ARRAY_ID)); // [1, 2]
    assertNode(yieldedValue.getFirstChild()).hasColorThat().isEqualTo(StandardColors.NUMBER); // 1
    assertNode(yieldedValue.getSecondChild()).hasColorThat().isEqualTo(StandardColors.NUMBER); // 2

    Node zero = yieldedValue.getNext();
    checkState(0 == zero.getDouble(), zero);
    assertNode(zero).hasColorThat().isEqualTo(StandardColors.NUMBER);
  }

  @Test
  public void testGeneratorForIn_withTypes() {
    Node case1Node =
        testAndReturnBodyForNumericGenerator(
            "for (var i in []) { yield 3; };",
            """
            var i;
            var GEN_FORIN$0$0;
            """,
            """
            if (GEN_CONTEXT$0.nextAddress == 1) {
              GEN_FORIN$0$0 = GEN_CONTEXT$0.forIn([]);
            }
            if (GEN_CONTEXT$0.nextAddress != 4) {
              if (!((i = GEN_FORIN$0$0.getNext()) != null)) {
                return GEN_CONTEXT$0.jumpTo(4);
              }
              return GEN_CONTEXT$0.yield(3, 2);
            }
            ;
            GEN_CONTEXT$0.jumpToEnd();
            """);

    // GEN_FORIN$0$0 = GEN_CONTEXT$0.forIn([]);
    Node assign = case1Node.getSecondChild().getFirstFirstChild();
    checkState(assign.isAssign(), assign);

    Color propertyIterator = StandardColors.TOP_OBJECT;
    assertThat(assign.getColor()).isEqualTo(propertyIterator);
    assertThat(assign.getFirstChild().getColor()).isEqualTo(propertyIterator);
    assertThat(assign.getSecondChild().getColor()).isEqualTo(propertyIterator);

    // if (!((i = GEN_FORIN$0$0.getNext()) != null)) {
    Node case2Node = case1Node.getNext();
    Node ifNode = case2Node.getSecondChild().getFirstChild();
    checkState(ifNode.isIf(), ifNode);
    Node ifCond = ifNode.getFirstChild();
    checkState(ifCond.isNot(), ifCond);
    assertNode(ifCond).hasColorThat().isEqualTo(StandardColors.BOOLEAN);
    Node ne = ifCond.getFirstChild();
    assertNode(ifCond).hasColorThat().isEqualTo(StandardColors.BOOLEAN);

    Node lhs = ne.getFirstChild(); // i = GEN_FORIN$0$0.getNext()
    assertNode(lhs)
        .hasColorThat()
        .hasAlternates(StandardColors.NULL_OR_VOID, StandardColors.STRING);
    assertNode(lhs)
        .hasColorThat()
        .hasAlternates(StandardColors.NULL_OR_VOID, StandardColors.STRING);
    assertNode(lhs)
        .hasColorThat()
        .hasAlternates(StandardColors.NULL_OR_VOID, StandardColors.STRING);
    Node getNextFn = lhs.getSecondChild().getFirstChild();
    assertNode(getNextFn).hasColorThat().isEqualTo(StandardColors.TOP_OBJECT);

    Node rhs = ne.getSecondChild();
    checkState(rhs.isNull(), rhs);
    assertNode(rhs).hasColorThat().isEqualTo(StandardColors.NULL_OR_VOID);

    // GEN_CONTEXT$0.jumpToEnd()
    Node case4Node = case2Node.getNext();
    Node jumpToEndCall = case4Node.getNext().getFirstChild();
    checkState(jumpToEndCall.isCall());

    Node jumpToEndFn = jumpToEndCall.getFirstChild();
    Node jscompGeneratorContext = jumpToEndFn.getFirstChild();

    // assertThat(jumpToEndCall.getJSType().toString()).isEqualTo("undefined");

    Color generatorContext =
        Iterables.getOnlyElement(
            CodeSubTree.findFirstNode(
                    getLastCompiler().getExternsRoot(),
                    (n) -> n.matchesQualifiedName("$jscomp.generator.Context"))
                .getColor()
                .getInstanceColors());
    assertThat(jscompGeneratorContext.getColor()).isEqualTo(generatorContext);
  }

  @Test
  public void testGeneratorTryCatch_withTypes() {
    Node case1Node =
        testAndReturnBodyForNumericGenerator(
            "try {yield 1;} catch (e) {}",
            "var e;",
            """
            if (GEN_CONTEXT$0.nextAddress == 1) {
              GEN_CONTEXT$0.setCatchFinallyBlocks(2);
              return GEN_CONTEXT$0.yield(1, 4);
            }
            if (GEN_CONTEXT$0.nextAddress != 2) {
              return GEN_CONTEXT$0.leaveTryBlock(0)
            }
            e = GEN_CONTEXT$0.enterCatchBlock();
            GEN_CONTEXT$0.jumpToEnd();
            """);
    Node case2Node = case1Node.getNext().getNext();

    // Test that "e = GEN_CONTEXT$0.enterCatchBlock();" has the unknown type
    Node eAssign = case2Node.getFirstChild();
    checkState(eAssign.isAssign(), eAssign);
    assertNode(eAssign).hasColorThat().isEqualTo(StandardColors.UNKNOWN);
    assertNode(eAssign.getFirstChild()).hasColorThat().isEqualTo(StandardColors.UNKNOWN);
    assertNode(eAssign.getSecondChild()).hasColorThat().isEqualTo(StandardColors.UNKNOWN);

    Node enterCatchBlockFn = eAssign.getSecondChild().getFirstChild();
    checkState(enterCatchBlockFn.isGetProp());
    // this is loosely typed, but it's okay because optimizations don't care about callee types.
    assertNode(enterCatchBlockFn).hasColorThat().isEqualTo(StandardColors.UNKNOWN);
  }

  @Test
  public void testGeneratorMultipleVars_withTypes() {
    Node firstStatementExprResult =
        testAndReturnBodyForNumericGenerator(
            "var x = 1, y = '2';", //
            """
            var x;
            var y;
            """,
            """
            x = 1;
            y = '2';
            GEN_CONTEXT$0.jumpToEnd();
            """);
    // x = 1
    final NodeSubject firstStatementSubject = assertNode(firstStatementExprResult).isExprResult();
    final NodeSubject assignASubject = firstStatementSubject.hasOneChildThat().isAssign();
    assignASubject.hasColorThat().isEqualTo(StandardColors.NUMBER);
    assignASubject.hasFirstChildThat().hasColorThat().isEqualTo(StandardColors.NUMBER);
    assignASubject.hasSecondChildThat().hasColorThat().isEqualTo(StandardColors.NUMBER);

    // y = '2';
    final NodeSubject secondStatementSubject =
        assertNode(firstStatementExprResult.getNext()).isExprResult();
    final NodeSubject assignBSubject = secondStatementSubject.hasOneChildThat().isAssign();
    assignBSubject.hasColorThat().isEqualTo(StandardColors.STRING);
    assignBSubject.hasFirstChildThat().hasColorThat().isEqualTo(StandardColors.STRING);
    assignBSubject.hasSecondChildThat().hasColorThat().isEqualTo(StandardColors.STRING);
  }

  /**
   * Tests that the given generator transpiles to the given body, and does some basic checks on the
   * transpiled generator.
   *
   * @return The first case statement in the switch inside the transpiled generator
   */
  private Node testAndReturnBodyForNumericGenerator(
      String beforeBody, String varDecls, String afterBody) {
    rewriteGeneratorBodyWithVarsAndReturnType(
        beforeBody, varDecls, afterBody, "!Generator<number>");

    Node transpiledGenerator = getLastCompiler().getJsRoot().getLastChild().getLastChild();
    Node program = getAndCheckGeneratorProgram(transpiledGenerator);

    Node programBlock = NodeUtil.getFunctionBody(program);
    return programBlock.getFirstChild();
  }

  /** Get the "program" function from a tranpsiled generator */
  private Node getAndCheckGeneratorProgram(Node genFunction) {
    Node returnNode = genFunction.getLastChild().getLastChild();
    Node callNode = returnNode.getFirstChild();
    checkState(callNode.isCall(), callNode);

    Node createGenerator = callNode.getFirstChild();
    return createGenerator.getNext().getNext();
  }

  @Test
  public void testBreakInNestedFinally() {
    rewriteGeneratorSwitchBodyWithVars(
        """
        OUTER: do {
          try {
            try {
              yield 1;
            } finally {
              break OUTER;
            }
          } finally {
            break OUTER;
          }
        } while (false);
        """,
        "",
        """
          GEN_CONTEXT$0.setFinallyBlock(5);
          GEN_CONTEXT$0.setFinallyBlock(7);
          return GEN_CONTEXT$0.yield(1, 7);
        case 7:
          GEN_CONTEXT$0.enterFinallyBlock(0, 5);
          GEN_CONTEXT$0.jumpThroughFinallyBlocks(0);
          GEN_CONTEXT$0.leaveFinallyBlock(5);
          break;
        case 5:
          GEN_CONTEXT$0.enterFinallyBlock();
          GEN_CONTEXT$0.jumpThroughFinallyBlocks(0);
          GEN_CONTEXT$0.leaveFinallyBlock(2);
          break;
        case 2:
          if (false) {
            GEN_CONTEXT$0.jumpTo(1);
            break;
          }
          GEN_CONTEXT$0.jumpToEnd();
        """);
  }

  @Test
  public void testFinallyReturnCancelledByOuterBreak() {
    rewriteGeneratorSwitchBodyWithVars(
        """
        OUT: do {
          try {
            yield 1234;
            return 'innerTry';
          } finally {
            break OUT;
          }
        } while(false);
        """,
        "",
        """
          GEN_CONTEXT$0.setFinallyBlock(5);
          return GEN_CONTEXT$0.yield(1234, 7);
        case 7:
          return GEN_CONTEXT$0.return("innerTry");
        case 5:
          GEN_CONTEXT$0.enterFinallyBlock();
          GEN_CONTEXT$0.jumpThroughFinallyBlocks(0);
          GEN_CONTEXT$0.leaveFinallyBlock(2);
          break;
        case 2:
          if (false) {
            GEN_CONTEXT$0.jumpTo(1);
            break;
          }
          GEN_CONTEXT$0.jumpToEnd();
        """);
  }

  @Test
  public void testFinallyReturnCancelledByOuterBreak_nested() {
    rewriteGeneratorSwitchBodyWithVars(
        """
        OUT: do {
          try {
            try {
              yield 1;
              return 'innerTry';
            } finally {
              break OUT;
            }
          } finally {
            yield 2;
          }
        } while (false);
        """,
        "",
        """
          GEN_CONTEXT$0.setFinallyBlock(5);
          GEN_CONTEXT$0.setFinallyBlock(7);
          return GEN_CONTEXT$0.yield(1, 9);
        case 9:
          return GEN_CONTEXT$0.return("innerTry");
        case 7:
          GEN_CONTEXT$0.enterFinallyBlock(0, 5);
          GEN_CONTEXT$0.jumpThroughFinallyBlocks(0);
          GEN_CONTEXT$0.leaveFinallyBlock(5);
          break;
        case 5:
          GEN_CONTEXT$0.enterFinallyBlock();
          return GEN_CONTEXT$0.yield(2, 10);
        case 10:
          GEN_CONTEXT$0.leaveFinallyBlock(2);
          break;
        case 2:
          if (false) {
            GEN_CONTEXT$0.jumpTo(1);
            break;
          }
          GEN_CONTEXT$0.jumpToEnd();
        """);
  }

  @Test
  public void testContinueInNestedFinallyInLoop() {
    rewriteGeneratorSwitchBodyWithVars(
        """
        OUT: do {
          try {
            yield 1;
          } finally {
            try {
              yield 2;
            } finally {
              continue OUT;
            }
          }
        } while (false);
        """,
        "",
        """
          GEN_CONTEXT$0.setFinallyBlock(5);
          return GEN_CONTEXT$0.yield(1, 5);
        case 5:
          GEN_CONTEXT$0.enterFinallyBlock();
          GEN_CONTEXT$0.setFinallyBlock(8);
          return GEN_CONTEXT$0.yield(2, 8);
        case 8:
          GEN_CONTEXT$0.enterFinallyBlock(0, 0, 1);
          GEN_CONTEXT$0.jumpThroughFinallyBlocks(2);
          GEN_CONTEXT$0.leaveFinallyBlock(9, 1);
          break;
        case 9:
          GEN_CONTEXT$0.leaveFinallyBlock(2);
          break;
        case 2:
          if (false) {
            GEN_CONTEXT$0.jumpTo(1);
            break;
          }
          GEN_CONTEXT$0.jumpToEnd();
        """);
  }
}
